import functools


class Solution(object):
    def removeBoxes(self, boxes):
        """
        :type boxes: List[int]
        :rtype: int
        """

        @functools.cache
        def search(i, j, k):
            n = j - i + 1
            base = boxes[i]
            num_base = 0
            right = i
            while right <= j:
                if boxes[right] == base:
                    num_base += 1
                    right += 1
                else:
                    break
            if right > j:
                max_count = (k + num_base) ** 2
            else:
                left = right
                max_count = search(right, j, 0) + (k + num_base) ** 2
                while right <= j:
                    if boxes[right] == base:
                        if boxes[right - 1] != base:
                            max_count = max(max_count, search(left, right - 1, 0) + search(right, j, k + num_base))
                    right += 1
            return max_count

        return search(0, len(boxes) - 1, 0)


data = Solution()
boxes = [1, 2, 2, 1, 1, 1, 2, 1, 1, 2, 1, 2, 1, 1, 2, 2, 1, 1, 2, 2, 1, 1, 1, 2, 2, 2, 2, 1, 2, 1, 1, 2, 2, 1, 2, 1, 2,
         2, 2, 2, 2, 1, 2, 1, 2, 2, 1, 1, 1, 2, 2, 1, 2, 1, 2, 2, 1, 2, 1, 1, 1, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 1,
         1, 1, 1, 2, 2, 2, 2, 2, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1, 1, 2, 1, 2, 2, 1]
print(data.removeBoxes(boxes))
boxes = [3, 2, 3, 4, 3]
print(data.removeBoxes(boxes))
# boxes = [2, 1, 2]
# print(data.removeBoxes(boxes))